import paramiko
from scp import SCPClient
import sys
import time
import threading
import traceback

class TransferInterruptedException(Exception):
    """传输中断异常"""
    pass

class SCPTool:
    # 传输状态常量
    STATUS_IDLE = "空闲"
    STATUS_TRANSFERRING = "正在传输"
    STATUS_INTERRUPTED = "传输中断"
    STATUS_COMPLETED = "传输完成"
    STATUS_FAILED = "传输失败"

    def __init__(self, hostname, port=22, username=None, password=None, key_file=None):
        self.hostname = hostname
        self.port = port
        self.username = username
        self.password = password
        self.key_file = key_file
        self.ssh = None
        self.scp = None
        self._file_times = {}
        self._last_print_time = {}  # 记录每个文件上次打印时间
        self.progress_info = ""  # 当前进度信息字符串
        self.print_interval = 1.0  # 打印间隔（秒）
        self._interrupt_flag = False  # 中断标志
        self._is_transferring = False  # 是否正在传输
        self.print_to_console = True  # 是否实时打印到控制台
        self.transfer_status = self.STATUS_IDLE  # 当前传输状态
        self.custom_progress_callback = None
        self.custom_progress_callback_extra_data = {}

        self.progress_raw = {
            "filename": "",
            "is_transferring": False,
            "status": self.STATUS_IDLE,
            "sent": 0,
            "size": 0,
            "speed": 0,
            "eta": 0,
            "percent": 0,
            "sent_str": "",
            "size_str": "",
            "speed_str": "",
            "eta_str": "",
            "percent_str": "",
        }

    def connect(self):
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            if self.key_file:
                pkey = paramiko.RSAKey.from_private_key_file(self.key_file)
                self.ssh.connect(hostname=self.hostname, port=self.port, username=self.username, pkey=pkey)
            else:
                self.ssh.connect(hostname=self.hostname, port=self.port, username=self.username, password=self.password)
            self.scp = SCPClient(self.ssh.get_transport(), progress=self._progress)
            if self.print_to_console:
                print("SCP 连接成功")
        except Exception as e:
            if self.print_to_console:
                print(f"SCP 连接失败: {e}")
            raise e

    def _human_size(self, bytes_size):
        """将字节数转换为人类可读的大小格式"""
        if bytes_size == 0:
            return "0B"
        
        units = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
        unit_index = 0
        size = float(bytes_size)
        
        while size >= 1024 and unit_index < len(units) - 1:
            size /= 1024
            unit_index += 1
            
        if unit_index == 0:
            return f"{int(size)}B"
        else:
            return f"{size:.2f}{units[unit_index]}"

    def _human_speed(self, bytes_per_sec):
        """将传输速率转换为人类可读的速度格式"""
        if bytes_per_sec == 0:
            return "0B/s"
            
        units = ['B/s', 'KB/s', 'MB/s', 'GB/s', 'TB/s']
        unit_index = 0
        speed = float(bytes_per_sec)
        
        while speed >= 1024 and unit_index < len(units) - 1:
            speed /= 1024
            unit_index += 1
            
        if unit_index == 0:
            return f"{int(speed)}B/s"
        else:
            return f"{speed:.2f}{units[unit_index]}"

    def _human_time(self, seconds):
        """时间友好化显示"""
        if seconds < 0:
            return "未知"
        if seconds < 60:
            return f"{int(seconds)}秒"
        if seconds < 3600:
            return f"{int(seconds//60)}分{int(seconds%60)}秒"
        return f"{int(seconds//3600)}时{int((seconds%3600)//60)}分"
    
    def set_custom_progress_callback(self, callback,extra_data=None):
        self.custom_progress_callback = callback
        self.custom_progress_callback_extra_data = extra_data
    
    def set_socketio(self,socketio):
        self.socketio = socketio



    def _progress(self, filename, size, sent):
        # 检查中断标志
        if self._interrupt_flag:
            self.transfer_status = self.STATUS_INTERRUPTED
            interrupt_msg = f"✗ 传输被中断: {filename}"
            self.progress_info = interrupt_msg
            self._is_transferring = False
            self.progress_raw["filename"] = filename
            self.progress_raw["is_transferring"] = False
            self.progress_raw["status"] = self.STATUS_INTERRUPTED
            if self.print_to_console:
                print(f"\n{interrupt_msg}")
            raise TransferInterruptedException("传输被用户中断")

            
        now = time.time()
        
        # 每个文件首次记录起始时间
        if filename not in self._file_times:
            self._file_times[filename] = now
            self._last_print_time[filename] = 0  # 初始化上次打印时间
            
        start_time = self._file_times[filename]
        elapsed = now - start_time


        if self.custom_progress_callback:
            try:
                should_print = (now - self._last_print_time[filename] >= self.print_interval) or (sent == size)
                if should_print:
                    self.custom_progress_callback(filename, size, sent, self.custom_progress_callback_extra_data)
                    self._last_print_time[filename] = now
            except Exception as e:
                raise e
            return True

        percent = float(sent) / float(size) * 100 if size != 0 else 100
        bar_len = 30
        filled_len = int(round(bar_len * percent / 100))
        bar = '=' * filled_len + '-' * (bar_len - filled_len)

        speed = sent / elapsed if elapsed > 0 else 0
        if speed > 0 and size > sent:
            eta = (size - sent) / speed
        else:
            eta = 0


        # 使用优化后的单位显示
        sent_str = self._human_size(sent)
        size_str = self._human_size(size)
        speed_str = self._human_speed(speed)
        eta_str = self._human_time(eta)

        self.progress_raw = {
            "filename": str(filename),
            "is_transferring": True,
            "status": self.STATUS_TRANSFERRING,
            "sent": sent,
            "size": size,
            "speed": speed,
            "eta": eta,
            "percent": percent,
            "sent_str": sent_str,
            "size_str": size_str,
            "speed_str": speed_str,
            "eta_str": eta_str,
            "percent_str": f"{percent:.1f}%",
        }

        # 获取文件名（不含路径）
        display_filename = str(filename).split('/')[-1].split('\\')[-1]
        if len(display_filename) > 20:
            display_filename = display_filename[:17] + "..."

        # 构建进度信息字符串
        progress_str = f'{display_filename} |{bar}| {percent:.1f}% [{sent_str}/{size_str}] {speed_str} 剩余:{eta_str}'
        
        # 更新成员变量（始终更新）
        self.progress_info = progress_str
        
        # 检查是否需要打印（1秒间隔或传输完成）
        should_print = (now - self._last_print_time[filename] >= self.print_interval) or (sent == size)
        
        # 根据标志决定是否打印到控制台
        if should_print and self.print_to_console:
            sys.stdout.write(f'\r{progress_str}')
            sys.stdout.flush()
            self._last_print_time[filename] = now
        
        if sent == size:
            complete_msg = f"✓ 传输完成: {display_filename}"
            self.progress_info = complete_msg
            self._is_transferring = False
            self.progress_raw["filename"] = filename
            self.progress_raw["is_transferring"] = False
            self.progress_raw["status"] = self.STATUS_COMPLETED
            if self.print_to_console:
                print()  # 传输完成换行
            # 清理记录
            self._file_times.pop(filename, None)
            self._last_print_time.pop(filename, None)

    def interrupt_transfer(self):
        """中断当前传输"""
        if self._is_transferring:
            self._interrupt_flag = True
            if self.print_to_console:
                print("\n正在中断传输...")
        else:
            if self.print_to_console:
                print("当前没有正在进行的传输")

    def _reset_interrupt_flag(self):
        """重置中断标志"""
        self._interrupt_flag = False

    def get_progress_info(self):
        """获取当前进度信息"""
        return self.progress_info
    
    def get_progress_raw(self):
        """获取当前进度信息"""
        return self.progress_raw

    def get_transfer_status(self):
        """获取当前传输状态"""
        return self.transfer_status

    def set_print_interval(self, interval):
        """设置打印间隔（秒）"""
        self.print_interval = max(0.1, interval)  # 最小0.1秒

    def set_print_to_console(self, enabled):
        """设置是否打印到控制台"""
        self.print_to_console = enabled

    def is_transferring(self):
        """检查是否正在传输"""
        return self._is_transferring

    def upload(self, local_path, remote_path, recursive=False):
        self._is_transferring = True
        self._reset_interrupt_flag()
        self.transfer_status = self.STATUS_TRANSFERRING
        
        try:
            if self.print_to_console:
                print(f"开始上传: {local_path} -> {remote_path}")
            self.progress_info = f"正在上传: {local_path}"
            self.progress_raw["filename"] = local_path
            self.progress_raw["is_transferring"] = True
            self.progress_raw["status"] = self.STATUS_TRANSFERRING

            
            self.scp.put(local_path, remote_path, recursive=recursive)
            
            self.transfer_status = self.STATUS_COMPLETED
            if self.print_to_console:
                print(f"✓ 上传完成: {local_path}")
            
        except TransferInterruptedException:
            self.transfer_status = self.STATUS_INTERRUPTED
            if self.print_to_console:
                print(f"✗ 上传被中断: {local_path}")
            self.progress_info = f"✗ 上传被中断: {local_path}"
            self.progress_raw["filename"] = local_path
            self.progress_raw["is_transferring"] = False
            self.progress_raw["status"] = self.STATUS_INTERRUPTED

            raise
            
        except Exception as e:
            self.transfer_status = self.STATUS_FAILED
            if self.print_to_console:
                print(f"✗ 上传失败: {e}")
            self.progress_info = f"✗ 上传失败: {local_path}"
            self.progress_raw["filename"] = local_path
            self.progress_raw["is_transferring"] = False
            self.progress_raw["status"] = self.STATUS_FAILED
            raise
            
        finally:
            self._is_transferring = False
            # 清理文件记录
            self._file_times.clear()
            self._last_print_time.clear()

    def download(self, remote_path, local_path, recursive=False):
        self._is_transferring = True
        self._reset_interrupt_flag()
        self.transfer_status = self.STATUS_TRANSFERRING
        
        try:
            if self.print_to_console:
                print(f"开始下载: {remote_path} -> {local_path}")
            self.progress_info = f"正在下载: {remote_path}"
            self.progress_raw["filename"] = remote_path
            self.progress_raw["is_transferring"] = True
            self.progress_raw["status"] = self.STATUS_TRANSFERRING
            
            self.scp.get(remote_path, local_path, recursive=recursive)
            
            self.transfer_status = self.STATUS_COMPLETED
            if self.print_to_console:
                print(f"✓ 下载完成: {remote_path}")
            
        except TransferInterruptedException:
            self.transfer_status = self.STATUS_INTERRUPTED
            if self.print_to_console:
                print(f"✗ 下载被中断: {remote_path}")
            self.progress_info = f"✗ 下载被中断: {remote_path}"
            self.progress_raw["filename"] = remote_path
            self.progress_raw["is_transferring"] = False
            self.progress_raw["status"] = self.STATUS_INTERRUPTED
            raise
            
        except Exception as e:
            self.transfer_status = self.STATUS_FAILED
            if self.print_to_console:
                print(f"✗ 下载失败: {e}")
            self.progress_info = f"✗ 下载失败: {remote_path}"
            self.progress_raw["filename"] = remote_path
            self.progress_raw["is_transferring"] = False
            self.progress_raw["status"] = self.STATUS_FAILED
            # trackback
            traceback.print_exc()


            raise
            
        finally:
            self._is_transferring = False
            # 清理文件记录
            self._file_times.clear()
            self._last_print_time.clear()

    def close(self):
        # 如果正在传输，先中断
        if self._is_transferring:
            self.interrupt_transfer()
            
        if self.scp:
            self.scp.close()
        if self.ssh:
            self.ssh.close()
        
        self.transfer_status = self.STATUS_IDLE
        self.progress_info = "连接已关闭"
        self.progress_raw["filename"] = ""
        self.progress_raw["is_transferring"] = False
        self.progress_raw["status"] = self.STATUS_IDLE
        if self.print_to_console:
            print("连接已关闭")


# 在另一个线程中启动下载
def download_task():
    try:
        scp_tool.download("/mnt/data/sim_control.tgz", "C:\\Users\\cheng.chen\\Downloads\\test\\",True)
    except TransferInterruptedException:
        print("下载已被中断")
    except Exception as e:
        print(f"下载出错: {e}")
        # trackback
        traceback.print_exc()

def upload_task():
    try:
        scp_tool.upload("C:\\Users\\cheng.chen\\Downloads\\test\\sim_control.tgz", "/tmp/")
    except TransferInterruptedException:
        print("上传已被中断")
    except Exception as e:
        print(f"上传出错: {e}")
        # trackback
        traceback.print_exc()



def monitor_progress():
    while True:
        status = scp_tool.get_transfer_status()
        progress = scp_tool.get_progress_info()
        print(f"状态: {status} | 进度: {progress}")
        
        if status in [scp_tool.STATUS_COMPLETED, scp_tool.STATUS_FAILED, scp_tool.STATUS_INTERRUPTED]:
            break
        time.sleep(1)    

if __name__ == "__main__":
    scp_tool = SCPTool(hostname="172.24.9.213", username="mm", password="mm")
    
    # scp_tool.upload("local_file.txt", "/home/mm/")

    # 关闭控制台打印（静默模式）
    scp_tool.set_print_to_console(True)
    # 进度打印的时间间隔
    scp_tool.set_print_interval(1)
   
    scp_tool.connect()


    # 启动下载线程
    # download_thread = threading.Thread(target=download_task)
    upload_thread = threading.Thread(target=upload_task)

    # monitor_thread = threading.Thread(target=monitor_progress)

    # 下载线程守护
    # download_thread.daemon = True

    # download_thread.start()
    upload_thread.start()

    # monitor_thread.start()


    # 等待5秒后中断下载
    # time.sleep(5)
    # print("\n================================")
    # print("//////",scp_tool.get_progress_info(),'/////')
    # time.sleep(0.2)
    # scp_tool.interrupt_transfer()

    # 线程分离
    # download_thread.daemon = True
    # monitor_thread.join()


    # scp_tool.close()


#     帮我基于html和layui设计一个前端界面，将下面的功能包装为一个类，通过js生成html。点击一个“下载”按钮后，在弹出层中渲染以下功能：

# 1. 有一个单选按钮，分别是“scp直传”、“服务器中转”
# 2. 若点击“scp直传”按钮，下面将显示远程主机ip、远程主机上待下载的文件路径（可编辑）、本地主机ip(可编辑）、本地主机接受路径（可编辑或者弹出路径选中器）
# 3. 若点击“服务器中转”，